package com.liuqi.openai.chain.result;

import com.liuqi.openai.data.message.AiMessage;
import com.liuqi.openai.data.message.ChatMessage;
import com.liuqi.openai.data.message.ToolResultMessage;
import com.liuqi.openai.data.message.UserMessage;
import com.liuqi.openai.model.output.TokenUsage;
import java.util.LinkedList;
import java.util.List;
import lombok.Builder;
import lombok.Getter;
import lombok.ToString;

/**
 * 对话链中，一轮对话的结果集。
 * <p>
 * 如果有工具调用，那么就会有多条消息，这些消息都会记录到结果集中
 *
 * @author liuqi
 * @date 2025/7/20
 **/
@ToString
@Getter
public class ChainResult {

    /**
     * 对话 开始时间
     */
    private Long beginTime;

    /**
     * 对话 结束时间
     */
    private Long endTime;

    /**
     * 用户对话消息
     * <p>
     * 用户对话消息可能会格式化成其他模板，存储的不一定是原始数据
     * </p>
     */
    private UserMessage userMessage;

    /**
     * AI响应消息 或者是 工具执行结果
     */
    private List<AssistantMessage> assistantMessages;

    /**
     * 设置开始时间
     */
    public void begin() {
        if (beginTime != null) {
            throw new IllegalStateException("beginTime already set");
        }
        this.beginTime = System.currentTimeMillis();
    }

    /**
     * 设置结束时间
     */
    public void end() {
        if (endTime != null) {
            throw new IllegalStateException("endTime already set");
        }
        this.endTime = System.currentTimeMillis();
    }

    /**
     * 设置用户消息
     *
     * @param userMessage
     */
    public void setUserMessage(UserMessage userMessage) {
        if (this.userMessage != null) {
            throw new IllegalStateException("userMessage already set");
        }
        this.userMessage = userMessage;
    }

    /**
     * 添加AI消息
     *
     */
    public void addAssistantMessage(Long startTime,
                                     Long endTime,
                                     AiMessage aiMessage,
                                     Float thinkDuration,
                                     TokenUsage tokenUsage) {
        this.addAssistantMessage(
                AssistantMessage.builder()
                .startTime(startTime)
                .endTime(endTime)
                .assistantMessage(aiMessage)
                .thinkDuration(thinkDuration)
                .tokenUsage(tokenUsage)
                .build()
        );
    }

    /**
     * 添加工具执行结果消息
     *
     */
    public void addToolResultMessage(Long startTime,
                                    Long endTime,
                                    ToolResultMessage toolResultMessage) {
        this.addAssistantMessage(
                AssistantMessage.builder()
                        .startTime(startTime)
                        .endTime(endTime)
                        .assistantMessage(toolResultMessage)
                        .build()
        );
    }

    /**
     * 添加AI消息
     *
     * @param assistantMessage
     */
    public void addAssistantMessage(AssistantMessage assistantMessage) {
        if (this.assistantMessages == null) {
            this.assistantMessages = new LinkedList<>();
        }
        this.assistantMessages.add(assistantMessage);
    }

    /**
     * AI 响应消息
     */
    @Builder
    @Getter
    @ToString
    public static class AssistantMessage {
        /**
         * 开始时间
         */
        private Long startTime;

        /**
         * 结束对话时间
         */
        private Long endTime;

        /**
         * AI响应消息 或者是 工具执行结果
         *
         * @see AiMessage
         * @see ToolResultMessage
         */
        private ChatMessage assistantMessage;

        /**
         * 思考内容耗时
         *
         * 如果有思考内容，则需要设置深度思考的耗时
         */
        private Float thinkDuration;

        /**
         * tokens计量
         */
        private TokenUsage tokenUsage;

    }
}
